// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/exceptions.h>
#include <pollux/exec/window_function.h>
#include <pollux/expression/function_signature.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::functions::window {
    namespace {
        class RowNumberFunction : public exec::WindowFunction {
        public:
            explicit RowNumberFunction(const TypePtr &resultType)
                : WindowFunction(resultType, nullptr, nullptr) {
            }

            void resetPartition(const exec::WindowPartition * /*partition*/) override {
                rowNumber_ = 1;
            }

            void apply(
                const BufferPtr &peerGroupStarts,
                const BufferPtr & /*peerGroupEnds*/,
                const BufferPtr & /*frameStarts*/,
                const BufferPtr & /*frameEnds*/,
                const SelectivityVector &validRows,
                vector_size_t resultOffset,
                const VectorPtr &result) override {
                if (resultType_->is_integer()) {
                    applyTyped<int32_t>(peerGroupStarts, resultOffset, result);
                } else {
                    applyTyped<int64_t>(peerGroupStarts, resultOffset, result);
                }
            }

            template<typename T>
            void applyTyped(
                const BufferPtr &peerGroupStarts,
                vector_size_t resultOffset,
                const VectorPtr &result) {
                int numRows = peerGroupStarts->size() / sizeof(vector_size_t);
                auto *rawValues = result->as_flat_vector<T>()->mutableRawValues();
                for (int i = 0; i < numRows; i++) {
                    rawValues[resultOffset + i] = rowNumber_++;
                }
            }

        private:
            int64_t rowNumber_ = 1;
        };
    } // namespace

    void registerRowNumber(const std::string &name, TypeKind resultTypeKind) {
        std::vector<exec::FunctionSignaturePtr> signatures{
            exec::FunctionSignatureBuilder()
            .returnType(mapTypeKindToName(resultTypeKind))
            .build(),
        };

        exec::registerWindowFunction(
            name,
            std::move(signatures),
            {exec::WindowFunction::ProcessMode::kRows, false},
            [name](
        const std::vector<exec::WindowFunctionArg> & /*args*/,
        const TypePtr &resultType,
        bool /*ignoreNulls*/,
        pollux::memory::MemoryPool * /*pool*/,
        HashStringAllocator * /*stringAllocator*/,
        const core::QueryConfig & /*queryConfig*/)
        -> std::unique_ptr<exec::WindowFunction> {
                return std::make_unique<RowNumberFunction>(resultType);
            });
    }

    void registerRowNumberInteger(const std::string &name) {
        registerRowNumber(name, TypeKind::INTEGER);
    }

    void registerRowNumberBigint(const std::string &name) {
        registerRowNumber(name, TypeKind::BIGINT);
    }
} // namespace kumo::pollux::functions::window
